# -*- coding: utf-8 -*-
# @Time    : 2025/8/5 15:18
# @Author  : Zhaoxu
# @File    : together.py
# @Software: PyCharm

# -*- coding:utf-8 -*-
"""
Realtime Voice Chat  (Mic → XF ASR → GPT → XF TTS → PyAudio 播放)
-----------------------------------------------------------------
依赖:
    pip install openai websocket-client pyaudio
    # Windows 安装 PyAudio 可先:
    #   pip install pipwin
    #   pipwin install pyaudio
-----------------------------------------------------------------
用法:
    1. 麦克风说话 (最长 10 秒；静音 2.5 秒自动结束)
    2. 讯飞实时语音转写得到文本
    3. 交给 GPT 生成回答 (DeepSeek Chat 示例)
    4. GPT 答案边生成、边切块、边走讯飞 TTS 合成
    5. 每帧音频立即写入 PyAudio, 实时播放
"""
import queue
import threading
import time
import os
import base64
import json
import hashlib
import hmac
from urllib.parse import urlencode, quote
from datetime import datetime
from wsgiref.handlers import format_date_time

import websocket
import pyaudio
from websocket import create_connection
from openai import OpenAI

# ========= 连接池管理 ========= #
class ConnectionPool:
    """连接池管理类"""
    
    def __init__(self):
        self.asr_connection = None
        self.tts_connection = None
        self.pyaudio_instance = None
        self.is_initialized = False
        
    def initialize_connections(self, asr_appid, asr_apikey, tts_appid, tts_apikey, tts_apisecret):
        """初始化所有连接"""
        print("🔗 正在初始化连接池...")
        
        # 初始化PyAudio
        self.pyaudio_instance = pyaudio.PyAudio()
        
        # 初始化ASR连接
        self.asr_connection = XFRealtimeASR(asr_appid, asr_apikey)
        self.asr_connection.connect()
        
        # 初始化TTS连接（这里只是准备参数，实际连接在需要时建立）
        self.tts_appid = tts_appid
        self.tts_apikey = tts_apikey
        self.tts_apisecret = tts_apisecret
        
        self.is_initialized = True
        print("✅ 连接池初始化完成")
        
    def get_asr_connection(self):
        """获取ASR连接"""
        if not self.is_initialized:
            raise Exception("连接池未初始化")
        return self.asr_connection
        
    def get_tts_params(self):
        """获取TTS参数"""
        if not self.is_initialized:
            raise Exception("连接池未初始化")
        return self.tts_appid, self.tts_apikey, self.tts_apisecret
        
    def close_all(self):
        """关闭所有连接"""
        print("🔌 正在关闭所有连接...")
        
        if self.asr_connection:
            self.asr_connection.close()
            
        if self.pyaudio_instance:
            self.pyaudio_instance.terminate()
            
        print("✅ 所有连接已关闭")

# ========= GPT → 队列 ========= #
def gpt_stream_to_queue(q: queue.Queue, user_prompt: str, conversation_history=None):
    """
    GPT 流式输出 → 按规则切块 → 入队
    支持多轮对话历史
    """
    if conversation_history is None:
        conversation_history = []
        
    client = OpenAI(
        api_key="sk-BPVBVZQbXHrUH6b540DbE34b4c71479f9dD4106dCc29C021",#sk-9bc1ce71791b49f083e625f2f2788b2d",
        base_url="https://api.pumpkinaigc.online/v1"
    )
    
    # 构建消息列表，包含历史对话
    messages = [
        {"role": "system", "content": "You are chatting with me ,your name is xiaoxi only the user use chinese speaking ,your name is 小西,"}
    ]
    
    # 添加历史对话（最多保留最近3轮）
    if len(conversation_history) > 6:  # 3轮对话 = 6条消息
        conversation_history = conversation_history[-6:]
    
    messages.extend(conversation_history)
    messages.append({"role": "user", "content": user_prompt})
    
    resp = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=messages,
        stream=True
    )

    buffer = ""
    full_response = ""
    for part in resp:
        delta = part.choices[0].delta
        if not getattr(delta, "content", None):
            continue
        token = delta.content
        print(token, end="", flush=True)          # 在终端同步显示
        buffer += token
        full_response += token
        if len(buffer) >= 80 or buffer.endswith(('.', '!', '?', '。', '！', '？')):
            q.put(buffer)
            buffer = ""
    if buffer.strip():
        q.put(buffer)
    q.put(None)     # 结束信号
    
    # 返回完整的AI回答，用于更新对话历史
    return full_response.strip()


# ========= 讯飞 TTS websocket ========= #
class Ws_Param:
    def __init__(self, appid, apikey, apisecret, text):
        self.APPID = appid
        self.APIKey = apikey
        self.APISecret = apisecret
        self.Text = text
        self.CommonArgs = {"app_id": self.APPID}
        self.BusinessArgs = {"aue": "raw",
                             "auf": "audio/L16;rate=16000",
                             "vcn": "x4_yezi",
                             "tte": "utf8"}
        self.Data = {"status": 2,
                     "text": str(base64.b64encode(self.Text.encode('utf-8')), "utf8")}

    def create_url(self):
        url = "wss://tts-api.xfyun.cn/v2/tts"
        now = datetime.now()
        date = format_date_time(time.mktime(now.timetuple()))
        sig_origin = f"host: ws-api.xfyun.cn\n" \
                     f"date: {date}\n" \
                     f"GET /v2/tts HTTP/1.1"
        sign = hmac.new(self.APISecret.encode(), sig_origin.encode(),
                        digestmod=hashlib.sha256).digest()
        sign = base64.b64encode(sign).decode()
        auth_origin = f'api_key="{self.APIKey}", algorithm="hmac-sha256", ' \
                      f'headers="host date request-line", signature="{sign}"'
        authorization = base64.b64encode(auth_origin.encode()).decode()
        return url + "?" + urlencode({"authorization": authorization,
                                      "date": date,
                                      "host": "ws-api.xfyun.cn"})


def tts_stream_and_play(appid, apikey, apisecret, text, pyaudio_instance):
    """
    对 text 进行 TTS；收到音频帧立即写 PyAudio，阻塞直到播放完
    使用连接池中的PyAudio实例
    """
    if not text.strip():
        return
        
    # 使用连接池中的PyAudio实例
    stream = pyaudio_instance.open(format=pyaudio.paInt16,
                                  channels=1,
                                  rate=16000,
                                  output=True)

    done_flag = threading.Event()

    def on_message(ws, message):
        msg = json.loads(message)
        audio_bytes = base64.b64decode(msg["data"]["audio"])
        stream.write(audio_bytes)                  # 实时播放
        if msg["data"]["status"] == 2:
            # TTS合成完成，设置完成标志但不主动关闭连接
            done_flag.set()

    def on_error(ws, error):
        print(f"❌ TTS错误: {error}")

    def on_close(ws, close_status_code, close_msg):
        # 只在连接关闭时清理音频资源
        stream.stop_stream()
        stream.close()
        done_flag.set()

    def on_open(ws):
        def run(*_):
            req = {"common": wsParam.CommonArgs,
                   "business": wsParam.BusinessArgs,
                   "data": wsParam.Data}
            ws.send(json.dumps(req))
        threading.Thread(target=run).start()

    wsParam = Ws_Param(appid, apikey, apisecret, text)
    websocket.enableTrace(False)
    ws = websocket.WebSocketApp(wsParam.create_url(),
                                on_open=on_open,
                                on_message=on_message,
                                on_error=on_error,
                                on_close=on_close)
    ws.run_forever(sslopt={"cert_reqs": 0})
    done_flag.wait()           # 等播放完毕


def consumer_tts_play(q: queue.Queue, appid, apikey, apisecret, pyaudio_instance):
    """
    消费队列并播放TTS
    返回完整的AI回答文本
    使用连接池中的PyAudio实例
    """
    idx = 1
    full_response = ""
    while True:
        txt = q.get()
        if txt is None:
            break
        print(f"\n\n>>> GPT 第 {idx} 段（{len(txt)} 字）开始 TTS 并播放 …")
        tts_stream_and_play(appid, apikey, apisecret, txt, pyaudio_instance)
        full_response += txt
        idx += 1
    print("\n===== GPT 回答播放结束 =====")
    return full_response


# ========= 讯飞实时 ASR ========= #
def gen_signa(appid, apikey, ts):
    base = appid + ts
    md5 = hashlib.md5(base.encode()).hexdigest()
    signa = hmac.new(apikey.encode(), md5.encode(), digestmod='sha1').digest()
    return base64.b64encode(signa).decode()


class XFRealtimeASR:
    def __init__(self, appid, apikey, vad_eos_ms=1500):
        self.appid = appid
        self.apikey = apikey
        self.vad_eos_ms = vad_eos_ms
        self.ws = None
        self.handshake_done = threading.Event()
        self.closed = False
        self.end_tag = "{\"end\":true}"
        self.result_text = ""
        self.last_result_time = None
        self.record_start_time = None

    def connect(self):
        """建立ASR连接"""
        ts = str(int(time.time()))
        signa = gen_signa(self.appid, self.apikey, ts)
        url = f"ws://rtasr.xfyun.cn/v1/ws?appid={self.appid}&ts={ts}&signa={quote(signa)}&vad_eos={self.vad_eos_ms}"
        print("🔗 连接语音识别服务...")
        self.ws = create_connection(url)
        threading.Thread(target=self._recv_thread_func, daemon=True).start()

    def start_mic(self, conversation_history=None, connection_pool=None):
        """
        开始录音 + 实时识别，返回整句文本
        支持多轮对话历史，使用连接池
        """
        if not self.handshake_done.wait(10):
            print("❌ 握手失败")
            return "", ""
            
        CHUNK = 640
        FORMAT = pyaudio.paInt16
        CHANNELS = 1
        RATE = 16000
        
        # 使用连接池中的PyAudio实例
        if connection_pool and connection_pool.pyaudio_instance:
            pa = connection_pool.pyaudio_instance
        else:
            pa = pyaudio.PyAudio()
            
        stream = pa.open(format=FORMAT, channels=CHANNELS, rate=RATE,
                         input=True, frames_per_buffer=CHUNK)
        print("🎤 开始说话（支持中英文，静音2.5秒或10秒自动结束）")

        try:
            while not self.closed:
                data = stream.read(CHUNK, exception_on_overflow=False)
                self.ws.send(data, opcode=0x2)

                if self.record_start_time:
                    elapsed = time.time() - self.record_start_time
                    if elapsed >= 10 or (time.time() - self.last_result_time) > (self.vad_eos_ms / 1000):
                        print("✅ 语音识别结束")
                        q = queue.Queue()
                        # 直接调用GPT函数获取回答
                        ai_response = gpt_stream_to_queue(q, self.result_text, conversation_history)
                        
                        # 使用连接池中的TTS参数
                        if connection_pool:
                            tts_appid, tts_apikey, tts_apisecret = connection_pool.get_tts_params()
                            consumer_tts_play(q, tts_appid, tts_apikey, tts_apisecret, connection_pool.pyaudio_instance)
                        else:
                            # 备用方案
                            XF_APPID_TTS = "9646138c"
                            XF_APIKEY_TTS = "ea89fb89877aded124af01b32cb14324"
                            XF_SECRET_TTS = "MzlhYmQyMjI5ODYzYzU2ODM2YWMxYWNk"
                            consumer_tts_play(q, XF_APPID_TTS, XF_APIKEY_TTS, XF_SECRET_TTS, pa)
                            
                        self.ws.send(self.end_tag.encode())
                        break
                time.sleep(0.04)
        except KeyboardInterrupt:
            print("⏹️ 用户中断")
        finally:
            stream.stop_stream()
            stream.close()
            # 注意：这里不关闭PyAudio，因为它是连接池管理的
            time.sleep(0.2)
            return self.result_text, ai_response if 'ai_response' in locals() else ""

    def parse_data_text(self, data_str):
        try:
            data_dict = json.loads(data_str)
            words = []
            for rt in data_dict.get('cn', {}).get('st', {}).get('rt', []):
                for ws in rt['ws']:
                    for cw in ws['cw']:
                        words.append(cw['w'])
            return ''.join(words)
        except Exception as e:
            print(f"❌ 解析语音识别文本出错: {e}")
            return ""

    def _recv_thread_func(self):
        try:
            while self.ws.connected:
                result = self.ws.recv()
                if not result:
                    break
                j = json.loads(result)
                if j["action"] == "started":
                    self.handshake_done.set()
                elif j["action"] == "result":
                    txt = self.parse_data_text(j.get("data", ""))
                    if txt:
                        print(f"📝 识别: {txt}")
                        self.result_text = txt
                        self.last_result_time = time.time()
                        if self.record_start_time is None:
                            self.record_start_time = time.time()
                elif j["action"] == "error":
                    print(f"❌ 语音识别错误: {j}")
                    self.closed = True
                    break
        except Exception as e:
            if not self.closed:
                print(f"❌ 语音识别接收线程异常: {e}")

    def close(self):
        """关闭连接 - 在程序结束时调用"""
        self.closed = True
        try:
            if self.ws:
                self.ws.close()
                print("🔌 语音识别连接已关闭")
        except:
            pass


# ========= main ========= #
def main():
    """
    主函数 - 多轮对话系统（连接池版本）
    支持三轮连续对话，使用连接池管理
    """
    print("=" * 60)
    print("🤖 多轮对话系统 v2.0 (连接池版本)")
    print("💬 支持三轮连续对话")
    print("🎤 说话后自动识别并回答")
    print("🔊 AI回答实时语音播放")
    print("📝 对话历史会自动保存")
    print("🔗 使用连接池管理，提高效率")
    print("=" * 60)
    
    # ======= 讯飞账号配置 =======
    XF_APPID_ASR   = "9646138c"
    XF_APIKEY_ASR  = "82adf64eb83cf7b9ec6f4ae215943e0d"
    XF_APPID_TTS   = "9646138c"
    XF_APIKEY_TTS  = "ea89fb89877aded124af01b32cb14324"
    XF_SECRET_TTS  = "MzlhYmQyMjI5ODYzYzU2ODM2YWMxYWNk"
    
    # 初始化连接池
    connection_pool = ConnectionPool()
    
    try:
        # 初始化所有连接
        connection_pool.initialize_connections(
            XF_APPID_ASR, XF_APIKEY_ASR,
            XF_APPID_TTS, XF_APIKEY_TTS, XF_SECRET_TTS
        )
        
        # 初始化对话历史
        conversation_history = []
        
        # 三轮对话循环
        for round_num in range(1, 4):
            print(f"\n🔄 第 {round_num} 轮对话")
            print("-" * 40)
            
            # ---- Step1: 语音实时转写 ---- #
            print("🎤 开始语音识别...")
            asr = connection_pool.get_asr_connection()
            user_text, ai_response = asr.start_mic(conversation_history, connection_pool)
            print(f"\n📝 识别结果: {user_text}")

            if not user_text.strip():
                print("⚠️ 未识别到文本，跳过本轮")
                continue

            # 更新对话历史
            conversation_history.append({"role": "user", "content": user_text})
            if ai_response:
                conversation_history.append({"role": "assistant", "content": ai_response})
            
            print(f"✅ 第 {round_num} 轮对话完成")
            print(f"📊 当前对话历史: {len(conversation_history)} 条消息")
            
            # 如果不是最后一轮，等待一下
            if round_num < 3:
                print("⏳ 准备下一轮对话...")
                time.sleep(1)
        
        print("\n🎉 三轮对话全部完成！")
        
    except KeyboardInterrupt:
        print("\n👋 用户中断，程序退出")
    except Exception as e:
        print(f"\n❌ 程序出错: {e}")
        print("请检查网络连接和API配置")
    finally:
        # 程序结束时统一关闭所有连接
        connection_pool.close_all()


if __name__ == "__main__":
    main()
